/*
 * hmm.h
 *
 *  Created on: Jul 4, 2011
 *      Author: tqlong
 */

#ifndef HMM_H_
#define HMM_H_

#include <armadillo>
#include <vector>
#include <cassert>
#include <boost/math/special_functions/log1p.hpp>

class HiddenMarkovModel
{
public:
  typedef std::vector<int> seq_type;
  typedef std::vector<seq_type> vseq_type;
protected:
  typedef arma::mat mat_type;
  typedef arma::vec vec_type;
  typedef arma::Mat<arma::s32> trace_type;
public:
  HiddenMarkovModel(int n_state, int n_observation);

  void
  genSeq(int length, seq_type &sseq, seq_type &oseq);

  /** log of probability of a sequence */
  double
  probability(const seq_type &oseq);

  /** Maximum a priori state sequence (viterbi algorithm) */
  void
  map(const seq_type &oseq, seq_type &sseq);

  /** Training with EM algorithm (Baum-Welch) */
  void
  train(const vseq_type& data);

  /** Setters and getters */
  void
  setTransition(const mat_type &a);

  void
  setEmission(const mat_type &b);

  void
  setStart(const vec_type &pi);

  const mat_type&
  transition() const;

  const mat_type&
  emission() const;

  const vec_type&
  start() const;

  void
  dump();

  static void
  print_seq(const seq_type &sseq, const seq_type &oseq);
protected:
  void
  initialize();

  void
  toProbabilityMatrix(mat_type &a);

  template<typename V>
    int
    randomWheel(const V& v);

  void
  forward_pass();

  void
  backward_pass();

  void
  viterbi_pass();

  void
  initialize_accumulate();

  void
  accumulate();

  void
  reestimate();

  double
  logmul(double a, double b) const;

  double
  logdiv(double a, double b) const;

  double
  logadd(double a, double b) const;
protected:
  /** Number of states */
  int n_state_;
  /** Number of observations */
  int n_observation_;
  /** pi is the starting probability of each state */
  vec_type pi_, logpi_, accpi_;
  /** a is the transition probability matrix (a_ij = Prob(q_t+1 = s_j | q_t = s_i) */
  mat_type a_, loga_, acca_;
  /** b is the emission probability matrix (b_ik = Prob(o_t = v_k | q_t = s_i) */
  mat_type b_, logb_, accb_;
  /** current sequence (in processing) */
  seq_type seq_;
  /** Forward and backward probabilities */
  mat_type alpha_, beta_;
  /** Viterbi backward trace */
  trace_type trace_;
};

/** Pick a random position in a vector of probability, avoid position with zero probability */
template<typename V>
  int
  HiddenMarkovModel::randomWheel(const V& v)
  {
    double s = 0;

    for (unsigned int i = 0; i < v.n_elem; i++)
      s += v[i];
    if (s == 0) return -1; // no position is picked

    // pick a random number and track the cummulative sum
    double r = arma::randu() * s, cs = 0;
    for (unsigned int i = 0; i < v.n_elem; i++)
      {
        cs += v[i];
        if (r <= cs) return i;
      }

    return -1;
  }

typedef HiddenMarkovModel hmm;
#define LOGZERO log(0)

#endif /* HMM_H_ */
